#include "tensor.h"
#include "pmath.h"
#include <iostream>

// 获取data的元素数量
int data_size(Shape shape){
    return shape.x * shape.y * shape.z;
}

// 默认构造函数
Tensor::Tensor(){}

// 实现Tensor类中的方法
Tensor::Tensor(Shape shape, int seed){
    this->shape.x = shape.x;
    this->shape.y = shape.y;
    this->shape.z = shape.z;
    this->data = rand_arr(data_size(shape), seed);
}

Tensor::Tensor(Shape shape, float *data){
    this->shape.x = shape.x;
    this->shape.y = shape.y;
    this->shape.z = shape.z;
    this->data = data; 
}

// 获取张量维度
int get_dim(Shape shape){
    if (shape.x ==1 && shape.y == 1 && shape.z == 1){
        return 0;
    }else if (shape.y == 1 && shape.z == 1){
        return 1;
    }else if (shape.z == 1){
        return 2;
    }else{
        return 3;
    }
}
// 转置仅针对矩阵
// 矩阵转置,将原始矩阵直接转置
void Tensor::t(){
    if (get_dim(this->shape) != 2){
        std::string msg = "transpose error: the dim of shape is not 2!";
        std::cout << msg << std::endl;
    }else{
        float *new_data = new float[data_size(this->shape)];
        float *old_data = this->data;
        transpose(old_data, new_data, this->shape.y, this->shape.x);
        this->data = new_data;
        delete[] old_data;
    }
}

// 转置,将原始矩阵转置后赋值给新矩阵
Tensor *Tensor::T(){
    Tensor *new_tensor = new Tensor();
    if (get_dim(this->shape) != 2){
        std::string msg = "transpose error: the dim of shape is not 2!";
        std::cout << msg << std::endl;
        return new_tensor;
    }
    
    float *new_data = new float[data_size(this->shape)];
    float *old_data = this->data;
    transpose(old_data, new_data, this->shape.y, this->shape.x);
    new_tensor->shape.x = this->shape.y;
    new_tensor->shape.y = this->shape.x;
    new_tensor->data = new_data;
    return new_tensor;
}

void Tensor::print_shape(){
    int dim = get_dim(this->shape);
    if (dim == 0)
        std::cout << "(1,)" << dim << std::endl;
    if (dim == 1)
        std::cout << "(" << this->shape.x << ",)" << std::endl;
    if (dim == 2)
        std::cout << "(" << this->shape.y << "," << this->shape.x << ")" << std::endl;
    if (dim == 3)
        std::cout << "(" << this->shape.z << "," << this->shape.y << "," << this->shape.x << ")" << std::endl;
}

void Tensor::print_data(){
    int offset_z = this->shape.x * this->shape.y;
    int offset_y = this->shape.x;
    for (int i = 0; i < this->shape.z; i++){
        for (int j = 0; j < this->shape.y; j++){
            for (int k = 0; k < this->shape.x; k++){
                std::cout << this->data[i*offset_z + j*offset_y + k] << " ";
            }
            std::cout << std::endl;
        }
        if (this->shape.z != 1)
            std::cout << std::endl;   
    }     
}

// tensor创建
Tensor *tensorMalloc(Shape shape, float *data, int seed){
    if (data == NULL){
        return new Tensor(shape, seed);
    }else{
        return new Tensor(shape, data);
    }
}

void tensorFree(Tensor *a){
        delete[] a->data;
        delete a;
}

// 检查两个tensor的shape是否相等
int shape_equal(Tensor *a, Tensor *b){
    if(a->shape.x != b->shape.x || a->shape.y != b->shape.y || a->shape.z != b->shape.z){
        return -1;
    }
    return 1;
}

// 矩阵乘法形状检查
int shape_matmul_equal(Tensor *a, Tensor *b){
    if(a->shape.x != b->shape.y){
        return -1;
    }
    return 1;
}

Tensor *tensor_add(Tensor *a, Tensor *b){
    if(shape_equal(a,b) < 0){
        std::string tip = "The shape of two tensors are not equal!";
        std::cout << tip << std::endl;
        return (Tensor *)NULL;
    }
    
    int s = data_size(a->shape);
    float *arr = new float[s];
    add(a->data, b->data, arr, s, NULL);
    return new Tensor(a->shape, arr);
}

Tensor *tensor_sub(Tensor *a, Tensor *b){
    if(shape_equal(a,b) < 0){
        std::string tip = "The shape of two tensors are not equal!";
        std::cout << tip << std::endl;
        return (Tensor *)NULL;
    }

    int s = data_size(a->shape);
    float *arr = new float[s];
    add(a->data, b->data, arr, s, 1);
    return new Tensor(a->shape, arr);
}

// 矩阵数乘
Tensor *tensor_mul(float a, Tensor *b){
    int s = data_size(b->shape);
    float *arr = new float[s];
    lmul(a, b->data, arr, s);
    return new Tensor(b->shape, arr);
}

Tensor *tensor_matmul(Tensor *a, Tensor *b){
    if(shape_matmul_equal(a, b) < 0){
        std::string tip = "This function must run with dim > 1!";
        std::cout << tip << std::endl;
        return (Tensor *)NULL;
    }
    
    float *arr = new float[a->shape.y * b->shape.x];
    degmm(a->shape.y, b->shape.x, a->shape.x, a->data, b->data, arr);
    return new Tensor(Shape{b->shape.x, a->shape.y ,a->shape.z}, arr);
}

//abs 
Tensor *mabs(Tensor *a){
    int s = data_size(a->shape);
    float *arr = new float[s];
    tensor_abs(a->data, arr, s);
    return new Tensor(a->shape, arr);
}

//dim = 0时，将行进行复制，dim = 1时，将列进行复制
Tensor *cps(Tensor *a, int dim, int num){
    Tensor *t = new Tensor();
    int s = data_size(a->shape);
    if (dim == 0 && a->shape.y != 1){
        std::string tip = "The first dimension must be 1!";
        std::cout << tip << std::endl;
        return t;
    }
    if (dim == 1 && a->shape.x != 1){
        std::string tip = "The second dimension must be 1!";
        std::cout << tip << std::endl;
        return t;
    }

    sscp(a->data, t->data, s, dim, num);
    if (dim == 0){
        t->shape.x = a->shape.x;
        t->shape.y = num;
        t->shape.z = a->shape.z;
    }
    if (dim == 1){
        t->shape.x = num;
        t->shape.y = a->shape.y;
        t->shape.z = a->shape.z;
    }
    return t;
}

// dim = 0时,以列为单位，将每行的同列数据相加；dim = 1时，以行为单位
Tensor *tensor_mean(Tensor *a, int dim){
        float *arr;
        Shape new_shape;
        if(dim == 0){
            new_shape.x = a->shape.x;
            new_shape.y = 1;
            new_shape.z = a->shape.z;
            arr = new float[a->shape.x];
        }else if(dim == 1){
            new_shape.x = 1;
            new_shape.y = a->shape.y;
            new_shape.z = a->shape.z;
            arr = new float[a->shape.y];
        }
        mean(a->data, arr, a->shape.y, a->shape.x, dim);
        return new Tensor(new_shape, arr);
}

Tensor *tensor_sum(Tensor *a, int dim){
    float *arr;
    Shape new_shape;
    if(dim == 0){
        new_shape.x = a->shape.x;
        new_shape.y = 1;
        new_shape.z = a->shape.z;
        arr = new float[a->shape.x];  
    }else if(dim == 1){
        new_shape.x = 1;
        new_shape.y = a->shape.y;
        new_shape.z = a->shape.z;
        arr = new float[a->shape.y];
    }
    
    sum(a->data, arr, a->shape.y, a->shape.x, dim);
    return new Tensor(new_shape, arr);
}

// 生成全零张量
Tensor *zeros(Shape shape){
    int size = data_size(shape);
    float *data = new float[size];
    for (int i = 0; i < size; i++){
        data[i] = 0;
    }
    return new Tensor(shape, data);

}

// 生成全一张量
Tensor *ones(Shape shape){
    int size = data_size(shape);
    float *data = new float[size];
    for (int i = 0; i < size; i++){
        data[i] = 1;
    }
    return new Tensor(shape, data);
}

// 将张量复制到新的张量中
Tensor *tensorCopy(Tensor *x){
    int size = data_size(x->shape);
    float *new_data = new float[size];
    for (int i = 0; i < size; i++){
        new_data[i] = x->data[i];
    }
    return new Tensor(x->shape, new_data);
}

Tensor *flat(Tensor *x, int dim){
    if (dim <= 0){
        return (Tensor *)NULL;
    }
    
    int size = data_size(x->shape);
    float *data = new float[size];
    for (int i = 0; i < size; i++){
        data[i] = x->data[i];
    }
    Shape new_shape;
    if (dim == 1){
        new_shape.x = x->shape.x * x->shape.y;
        new_shape.y = x->shape.z;
        new_shape.z = 1;
    }

    if (dim == 2){
        new_shape.x = x->shape.x * x->shape.y * x->shape.z;
        new_shape.y = 1;
        new_shape.z = 1;
    }
    return new Tensor(new_shape, data);
}